package org.superb.huffman;

import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;

/**
 * @author Superb
 * @date 2020/9/18 - 20:41
 * @E_mail superb12580@163.com
 */
public class HuffmanCode {
    public static void main(String[] args) {

    }

    /**
     * 文件压缩
     * @param srcFile 目标文件
     * @param zipFile 压缩文件
     * @throws IOException
     */
    public static void zipFile(String srcFile, String zipFile) throws IOException {
        //读入目标文件返回byte[]
        byte[] bytes = srcFileBytes(srcFile);
        //传入压缩前的byte[]返回压缩后的byte[]
        byte[] bytesZip = huffmanZip(bytes);
        //传入压缩后的byte[]和生成的huffman编码表生成指定路径下的压缩文件
        bytesZipDstFile(bytesZip, huffmanCodes, zipFile);
    }

    /**
     * 文件解压
     * @param srcFile 解压路径
     * @param dstFile 压缩包路径
     * @throws IOException
     * @throws ClassNotFoundException
     */
    public static void decodFile(String srcFile, String dstFile) throws IOException, ClassNotFoundException {
        FileInputStream fis = new FileInputStream(dstFile);

        BufferedInputStream bis = new BufferedInputStream(fis);

        ObjectInputStream ois = new ObjectInputStream(bis);

        Object objBytes = ois.readObject();

        byte[] bytesZip = (byte[]) objBytes;

        Object objHuffmanCodes = ois.readObject();

        Map<Byte, String> huffmanCodes = (Map<Byte, String>) objHuffmanCodes;

        byte[] bytes = huffmanDecod(bytesZip, huffmanCodes);

        FileOutputStream fos = new FileOutputStream(srcFile);

        BufferedOutputStream bos = new BufferedOutputStream(fos);

        bos.write(bytes);

        ois.close();

        bis.close();

        bos.close();

    }

    /**
     * 数据压缩
     * @param bytes
     * @return byte[]
     */
    public static byte[] huffmanZip(byte[] bytes) {
        //统计每个字符出现的次数返回list
        List<Note> listNote = getNotes(bytes);
        //生成对应Huffman树返回根结点
        Note note = huffmanTree(listNote);
        //根据Huffman树生成对应Huffman编码表
        Map<Byte, String> huffmanCodes = getCodes(note);
        //压缩并返回
        return zip(bytes, huffmanCodes);
    }

    /**
     * 数据解压
     * @param bytes
     * @param huffmanCodes
     * @return
     */
    public static byte[] huffmanDecod(byte[] bytes, Map<Byte, String> huffmanCodes) {
        StringBuilder sb = new StringBuilder();
        //压缩后的byte[]转成二进制字符串
        for (int i = 0; i < bytes.length; i++) {
            sb.append(byteToBitString(i != bytes.length - 1, bytes[i]));
        }
        //反转Map表
        Map<String, Byte> huffmanCodex = new HashMap<>();
        for (Entry<Byte, String> entry : huffmanCodes.entrySet()) {
            huffmanCodex.put(entry.getValue(), entry.getKey());
        }
        List<Byte> listByte = new ArrayList<>();
        int flag;
        for (int i = 0; i < sb.length(); ) {
            flag = 1;
            while (true) {
                if (huffmanCodex.get(sb.substring(i, i + flag)) != null) {
                    listByte.add(huffmanCodex.get(sb.substring(i, i + flag)));
                    break;
                } else {
                    flag++;
                }
            }
            i += flag;
        }
        byte[] by = new byte[listByte.size()];
        for (int i = 0; i < by.length; i++) {
            by[i] = listByte.get(i);
        }
        return by;
    }



    /****************************以下内容为工具方法**********************************/


    /**
     * 传入压缩后的byte[]和生成的huffman编码表生成指定路径下的压缩文件
     * @param bytesZip
     * @return void
     * @throws IOException
     */
    private static void bytesZipDstFile(byte[] bytesZip,Map<Byte,String> huffmanCodes,String dstFile) throws IOException {
        //读取文件
        FileOutputStream fos = new FileOutputStream(dstFile);
        //加缓冲流
        BufferedOutputStream bos = new BufferedOutputStream(fos);
        //转对象流
        ObjectOutputStream oos = new ObjectOutputStream(bos);
        //将压缩后的数据写出到byte数组
        oos.writeObject(bytesZip);
        //刷新
        oos.flush();
        //将赫夫曼编码表也写出到指定路径下
        oos.writeObject(huffmanCodes);
        //刷新
        oos.flush();
        //关闭流
        oos.close();
        bos.close();
    }

    /**
     * 传入文件路径返回此文件的byte[] 压缩调用
     * @param srcFile
     * @return byte[]
     * @throws IOException
     */
    private static byte[] srcFileBytes(String srcFile) throws IOException {

        FileInputStream fis = new FileInputStream(srcFile);

        BufferedInputStream bis = new BufferedInputStream(fis);
        //创建适宜长度数组
        byte[] bytes = new byte[bis.available()];
        //读取目标路径文件到byte数组
        bis.read(bytes);
        //关闭流
        bis.close();

        return bytes;
    }


    /**
     * byte转二进制字符串
     *
     * @param flag 判断是否不是最后一个byte
     * @param b
     * @return
     */
    private static String byteToBitString(boolean flag, byte b) {
        int temp = b;
        //如果不是最后一个byte
        //正数需要补高位 负数不需要但补也没影响
        if (flag && b >= 0) {
            temp |= 256;
        }
        String str = Integer.toBinaryString(temp);
        if (!flag && b >= 0) {
            return str;
        }
        return str.substring(str.length() - 8);
    }

    /**
     * byte转二进制字符串
     *
     * @param bytes 原始byte[]
     * @return String
     */
    private static String byteToBitString(byte[] bytes) {
        int temp;
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < bytes.length; i++) {
            temp = bytes[i];
            temp |= 256;
            sb.append(Integer.toBinaryString(temp).substring(Integer.toBinaryString(temp).length() - 8));
        }
        return sb.toString();
    }

    /**
     * 传入压缩前byte[]和对应Huffman编码表生成压缩后的byte[]
     * @param bytes 对应的就是需要压缩的文件或字符串
     * @param huffmanCodes 赫夫曼编码表
     * @return byte[] 返回压缩后的byte数组
     */
    private static byte[] zip(byte[] bytes,Map<Byte,String> huffmanCodes) {
        StringBuilder sb = new StringBuilder();
        //遍历元byte数组
        for(byte b : bytes) {
            //去编码表中匹配每个字节，并将其编码连接一起到sb
            sb.append(huffmanCodes.get(b));
        }
        //此时sb里面存储的就是压缩后的二进制数据(全部由01组成)，但此时仍是用字符串进行存储，一个字符两个字节。
        //此时的sb字符串比元数据要大得多，所以不能用字符串存二进制数据，还是要以byte数组(sb中每八个字符转化为一个字节)的形式。
        //创建一个长度适宜的byte数组
        byte[] byteZip = new byte[(sb.length()+7)/8];//绝对够用
        int count = 0;
        //遍历 每八个字符一循环
        for (int i = 0; i < sb.length(); i+=8) {
            //如果是最后一次循环，字符串长度很可能不是八的整数倍，从而不足八位，额外处理。
            if(count==byteZip.length-1) {
                byteZip[count] = (byte)Integer.parseInt(sb.substring(i),2);
                break;
            }
            //如果不是最后一次循环 就每八个字符生成一个byte
            byteZip[count++] = (byte)Integer.parseInt(sb.substring(i,i+8),2);
        }
        //返回生成byte[] 此时已经是压缩完毕的byte数组，字节远小于压缩前。
        return byteZip;
    }


    //成员变量 各byte对应的字符串二进制码
    private static Map<Byte,String> huffmanCodes = new HashMap<>();
    /**
     * 递归生成各Byte对应的字符串二进制码
     * @param note 根结点
     * @param str 0代表往左子树遍历 1代表往右子树遍历
     * @param sb
     */
    private static void getCodes(Note note,String str,StringBuilder sb) {
        //额外定义sb2接收每次递归生成的上层编码sb，不直接对sb进行操作，从而不影响全局变量sb
        StringBuilder sb2 = new StringBuilder(sb);
        //每次递归深入下一层先加上每一层的编码(左0右1)
        sb2.append(str);
        //先判不为空 再继续
        if(note!=null) {
            //如果数据为空，说明不是叶子结点，不是byte字节结点
            if(note.data==null) {
                //继续左右递归遍历 往左走+0 往右走+1
                getCodes(note.left, "0",sb2);
                getCodes(note.right, "1",sb2);
            }
            //如果不为空，说明是叶子结点，已递归生成其赫夫曼编码
            else {
                //存到map中 key=结点byte值 value=对应赫夫曼编码
                huffmanCodes.put(note.data, sb2.toString());
            }
        }
    }
    /**
     * 赫夫曼树转赫夫曼编码表(重载)
     * 调用上述方法，并返回生成的赫夫曼编码表(每个字节byte对应一个编码)
     * @param note
     * @return
     */
    private static Map<Byte,String> getCodes(Note note) {
        getCodes(note,"",new StringBuilder());
        //返回每个字节byte对应的赫夫曼编码：0110010(由0和1排列组合形成)
        return huffmanCodes;
    }


    //成员变量 存每个byte出现的次数
    private static Map<Byte,Integer> map = new HashMap<>();
    /**
     * 通过Map统计各byte出现的次数生成Note对象返回List<Note>
     * @param bytes
     * @return List<Note>
     */
    private static List<Note> getNotes(byte[] bytes){
        //存储每个字节生成的Node结点对象
        List<Note> list = new ArrayList<>();
        //遍历byte数组
        for (byte b : bytes) {
            //如果map中还没有此字节，就设次数为1
            if(map.get(b)==null) {
                map.put(b, 1);
            }
            //如果map中已出现此字节，就+1
            else {
                map.put(b, map.get(b)+1);
            }
        }
        //遍历整个map
        for(Entry<Byte, Integer> entry : map.entrySet()) {
            //为map中每个的字节创建Note结点，并传入字节数据和权重（出现次数），
            list.add(new Note(entry.getKey(),entry.getValue()));
        }
        //返回，此时已包含所有结点对象
        return list;
    }

    /**
     * 传入集合生成赫夫曼树返回根结点
     * @param list
     * @return Note
     */
    private static Note huffmanTree(List<Note> list) {
        Note left = null;
        Note right = null;
        Note parent = null;
        //循环遍历List集合中的结点，只要不为空就继续
        while(list.size()>1) {
            //对所有结点进行排序，前面已实现Comparable接口，调用Collections工具类即可排序。
            Collections.sort(list);
            //0和1位置的结点权值最小，分别赋值给左右结点
            left = list.get(0);
            right = list.get(1);
            //并生成它们的父结点，(权值相加，字节数据为空，因为最后不是叶子结点)
            parent = new Note(null,left.weight+right.weight);
            //连接左右孩子
            parent.left = left;
            parent.right = right;
            //从集合中移除已经生成的结点
            list.remove(left);
            list.remove(right);
            //添加新的父类结点
            list.add(parent);
        }
        //遍历结束 返回根结点 此时已生成带权路径长度最优的赫夫曼树
        return parent;
    }

}

/**
 * Note类(一个字节对应一个Note对象)
 * 实现Comparable接口根据权重大小进行排序
 * @author Superb
 */
class Note implements Comparable<Note>{
    //存储的字节数据
    Byte data;
    //权重(该字节数据出现的次数)非常重要
    int weight;
    //左节点
    Note left;
    //右节点
    Note right;
    //构造器 每创建一个Node节点就要传入其字节数据以及权重
    public Note(Byte data, int weight) {
        super();
        this.data = data;
        this.weight = weight;
    }
    @Override
    public String toString() {
        return "Node [data=" + data + ", weight=" + weight + "]";
    }
    //重写compareTo方法
    @Override
    public int compareTo(Note o) {
        //根据权重属性值 升序排列
        return this.weight-o.weight;
    }
}
